package decision_tree.metric;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import decision_tree.dto.Sample;
import decision_tree.question.Question;

public class BinaryEntropyManager extends EntropyManager{
	
	@Override
	public Map<Integer, List<Sample>> splitData(List<Sample> sampleList, Question q) {
		
		List<Sample> positive = new ArrayList<Sample>(sampleList.size()); 
		List<Sample> negative = new ArrayList<Sample>(sampleList.size()); 

		for(Sample sample : sampleList) {
			int category = q.ask(sample);
			if(category == 1) {
				positive.add(sample);
			} else {
				negative.add(sample);
			}
		}
		Map<Integer, List<Sample>> map = new HashMap<Integer, List<Sample>>();
		map.put(1, positive);
		map.put(0, negative);

		return map;
	}
}
